'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: token_buffer_memory.py
* @Time: 2025/10/9
* @All Rights Reserve By Brtc
'''
from dataclasses import dataclass

from langchain_core.language_models import BaseLanguageModel
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, trim_messages, get_buffer_string
from sqlalchemy.sql.expression import desc

from internal.entity.conversation_entity import MessageStatus
from internal.model import Conversation, Message
from pkg.sqlachemy.sqlalchemy import SQLAlchemy


@dataclass
class TokenBufferMemory:
    """基于token计数的缓冲记忆组件"""
    db:SQLAlchemy
    conversation:Conversation # 会话模型
    model_instance:BaseLanguageModel # 记忆总结大语言模型

    def get_history_prompt_message(self, max_token_limit:int=2000, message_limit:int=10)->list[AnyMessage]:
        """根据传递的会话消息列表 + 消息条数 限制和过去指定内容的历史消息列表"""
        #1、判断会话模型是否存在，如果不存在则返回空列表
        if self.conversation is None:
            return []

        #2、查询该会话的消息列表, 并使用时间进行倒序，同时匹配答案不能为空， 匹配会话id ， 状态是正常且没有没 删除
        messages = self.db.session.query(Message).filter(
            Message.conversation_id == self.conversation.id,
            Message.answer != "",
            Message.is_deleted == False,
            Message.status == MessageStatus.NORMAL
        ).order_by(desc("created_at")).limit(message_limit).all()
        messages = list(reversed(messages))

        #3、将message 转换成Langchain 消息列表
        prompt_message = []

        for message in messages:
            prompt_message.extend([
                HumanMessage(content=message.query),
                AIMessage(content=message.answer)
            ])

        #4、调用Langchain的trim_message 函数实现剪切消息列表
        return trim_messages(
            message = prompt_message,
            max_tokens = max_token_limit,
            strategy="last",
        )


    def get_history_prompt_text(self, human_prefix:str = "Human", ai_prefix:str="AI", max_token_limit:int=2000, message_limit:int=10)->str:
        """根据传递的数据获取指定的会话历史消息提示文本(短期记忆文本形式)"""
        #1、根据传递的信息获取历史消息列表
        messages = self.get_history_prompt_message(max_token_limit, message_limit)
        #2、调用Langchain 集成的的get_buffer_string函数将消息列表 转换成文本
        return get_buffer_string(messages, human_prefix, ai_prefix)

